import torch
import warnings
warnings.filterwarnings("ignore", category=Warning)

dataset = torch.load("../data/AM_data.pt")

print(f"数据集dataset = dataset[0]之前的(dataset)内容:{dataset}")
print(f"数据集dataset = dataset[0]之前的(type)内容:{type(dataset)}")
dataset = dataset[0]
# 提取相关信息
# features = dataset['x']  # 节点特征
edges = dataset['edge_index']  # 边索引
labels = dataset['y']  # 节点标签

# 打印基本信息
print("数据集信息：")
print(f"节点特征 (features) 形状: {features.shape}")  # 打印特征矩阵形状
print(f"边索引 (edges) 形状: {edges.shape}")  # 打印边索引形状
print(f"节点标签 (labels) 形状: {labels.shape}")  # 打印标签向量形状
print(f"数据集dataset = dataset[0]之后的(dataset)内容:{dataset}")
print(f"数据集dataset = dataset[0]之后的(type)内容:{type(dataset)}")

# 打印前5个节点的特征
print("\n前5个节点的特征（features）:")
print(features[:5])

# 打印前5条边的索引
print("\n前5条边的索引（edges）:")
print(edges[:, :5])

# 打印前5个节点的标签
print("\n前5个节点的标签（labels）:")
print(labels[:5])

# 统计每个类别的节点数量
unique_labels, counts = torch.unique(labels, return_counts=True)
print("\n每个类别的节点数量：")
for label, count in zip(unique_labels, counts):
    print(f"类别 {label.item()}: {count.item()} 个节点")